import pretty_midi
from collections import defaultdict 
from enum import Enum

class InstrumentType(Enum):
  ACOUSTIC_PIANO = 1
  ELECTRIC_PIANO = 2
  OTHER_PIANO = 3
  PERCUSSION = 4
  ORGAN = 5
  ACOUSTIC_GUITAR = 6
  ELECTRIC_GUITAR = 7
  OTHER_GUITAR = 8
  BASS = 9
  VIOLIN = 10
  VIOLA = 11
  CELLO = 12
  OTHER_STRING = 13
  ENSEMBLE = 14
  BRASS = 15
  REED = 16
  PIPE = 17
  SYNTH_LEAD = 18
  SYNTH_PAD = 19
  SYNTH_EFFECTS = 20
  ETHNIC = 21
  PERCUSSIVE = 22
  SOUND_EFFECTS = 23



# A program number to instrument type mapping.
PROGRAM_TO_INSTRUMENT_CATEGORY = {
     0: InstrumentType.ACOUSTIC_PIANO,
     1: InstrumentType.ACOUSTIC_PIANO,
     2: InstrumentType.ACOUSTIC_PIANO,
     3: InstrumentType.ACOUSTIC_PIANO,
     4: InstrumentType.ELECTRIC_PIANO,
     5: InstrumentType.ELECTRIC_PIANO,
     6: InstrumentType.OTHER_PIANO,
     7: InstrumentType.OTHER_PIANO,
     8: InstrumentType.PERCUSSION,
     9: InstrumentType.PERCUSSION,
     10: InstrumentType.PERCUSSION,
     11: InstrumentType.PERCUSSION,
     12: InstrumentType.PERCUSSION,
     13: InstrumentType.PERCUSSION,
     14: InstrumentType.PERCUSSION,
     15: InstrumentType.PERCUSSION,
     16: InstrumentType.ORGAN,
     17: InstrumentType.ORGAN,
     18: InstrumentType.ORGAN,
     19: InstrumentType.ORGAN,
     20: InstrumentType.ORGAN,
     21: InstrumentType.ORGAN,
     22: InstrumentType.ORGAN,
     23: InstrumentType.ORGAN,
     24: InstrumentType.ACOUSTIC_GUITAR,
     25: InstrumentType.ACOUSTIC_GUITAR,
     26: InstrumentType.ELECTRIC_GUITAR,
     27: InstrumentType.ELECTRIC_GUITAR,
     28: InstrumentType.ELECTRIC_GUITAR,
     29: InstrumentType.OTHER_GUITAR,
     30: InstrumentType.OTHER_GUITAR,
     31: InstrumentType.OTHER_GUITAR,
     32: InstrumentType.BASS,
     33: InstrumentType.BASS,
     34: InstrumentType.BASS,
     35: InstrumentType.BASS,
     36: InstrumentType.BASS,
     37: InstrumentType.BASS,
     38: InstrumentType.BASS,
     39: InstrumentType.BASS,
     40: InstrumentType.VIOLIN,
     41: InstrumentType.VIOLA,
     42: InstrumentType.CELLO,
     43: InstrumentType.OTHER_STRING,
     44: InstrumentType.OTHER_STRING,
     45: InstrumentType.OTHER_STRING,
     46: InstrumentType.OTHER_STRING,
     47: InstrumentType.OTHER_STRING,
     48: InstrumentType.ENSEMBLE,
     49: InstrumentType.ENSEMBLE,
     50: InstrumentType.ENSEMBLE,
     51: InstrumentType.ENSEMBLE,
     52: InstrumentType.ENSEMBLE,
     53: InstrumentType.ENSEMBLE,
     54: InstrumentType.ENSEMBLE,
     55: InstrumentType.ENSEMBLE,
     56: InstrumentType.BRASS,
     57: InstrumentType.BRASS,
     58: InstrumentType.BRASS,
     59: InstrumentType.BRASS,
     60: InstrumentType.BRASS,
     61: InstrumentType.BRASS,
     62: InstrumentType.BRASS,
     63: InstrumentType.BRASS,
     64: InstrumentType.REED,
     65: InstrumentType.REED,
     66: InstrumentType.REED,
     67: InstrumentType.REED,
     68: InstrumentType.REED,
     69: InstrumentType.REED,
     70: InstrumentType.REED,
     71: InstrumentType.REED,
     72: InstrumentType.PIPE,
     73: InstrumentType.PIPE,
     74: InstrumentType.PIPE,
     75: InstrumentType.PIPE,
     76: InstrumentType.PIPE,
     77: InstrumentType.PIPE,
     78: InstrumentType.PIPE,
     79: InstrumentType.PIPE,
     80: InstrumentType.SYNTH_LEAD,
     81: InstrumentType.SYNTH_LEAD,
     82: InstrumentType.SYNTH_LEAD,
     83: InstrumentType.SYNTH_LEAD,
     84: InstrumentType.SYNTH_LEAD,
     85: InstrumentType.SYNTH_LEAD,
     86: InstrumentType.SYNTH_LEAD,
     87: InstrumentType.SYNTH_LEAD,
     88: InstrumentType.SYNTH_PAD,
     89: InstrumentType.SYNTH_PAD,
     90: InstrumentType.SYNTH_PAD,
     91: InstrumentType.SYNTH_PAD,
     92: InstrumentType.SYNTH_PAD,
     93: InstrumentType.SYNTH_PAD,
     94: InstrumentType.SYNTH_PAD,
     95: InstrumentType.SYNTH_PAD,
     96: InstrumentType.SYNTH_EFFECTS,
     97: InstrumentType.SYNTH_EFFECTS,
     98: InstrumentType.SYNTH_EFFECTS,
     99: InstrumentType.SYNTH_EFFECTS,
     100: InstrumentType.SYNTH_EFFECTS,
     101: InstrumentType.SYNTH_EFFECTS,
     102: InstrumentType.SYNTH_EFFECTS,
     103: InstrumentType.SYNTH_EFFECTS,
     104: InstrumentType.ETHNIC,
     105: InstrumentType.ETHNIC,
     106: InstrumentType.ETHNIC,
     107: InstrumentType.ETHNIC,
     108: InstrumentType.ETHNIC,
     109: InstrumentType.ETHNIC,
     110: InstrumentType.ETHNIC,
     111: InstrumentType.ETHNIC,
     112: InstrumentType.PERCUSSIVE,
     113: InstrumentType.PERCUSSIVE,
     114: InstrumentType.PERCUSSIVE,
     115: InstrumentType.PERCUSSIVE,
     116: InstrumentType.PERCUSSIVE,
     117: InstrumentType.PERCUSSIVE,
     118: InstrumentType.PERCUSSIVE,
     119: InstrumentType.PERCUSSIVE,
     120: InstrumentType.SOUND_EFFECTS,
     121: InstrumentType.SOUND_EFFECTS,
     122: InstrumentType.SOUND_EFFECTS,
     123: InstrumentType.SOUND_EFFECTS,
     124: InstrumentType.SOUND_EFFECTS,
     125: InstrumentType.SOUND_EFFECTS,
     126: InstrumentType.SOUND_EFFECTS,
     127: InstrumentType.SOUND_EFFECTS,
}

def file_contains_only_allowed_programs(file_path, allowed_programs=None, allow_drums=True, allow_drum_programs=False):
  """
  allowed_instruments: An instrument category type to the number of allowed instrument channels mapping.
  """
  try:
    midi_data = pretty_midi.PrettyMIDI(str(file_path))
  except:
    return None
 
  for instrument in midi_data.instruments:
    if allowed_programs is not None:
      if instrument.program not in allowed_programs:
        return None
    if not allow_drums:
      if instrument.is_drum:
        return None
    elif not allow_drum_programs:
      if instrument.is_drum and instrument.program != 0:
        return None
  return midi_data